from rmasm.utils import *
from rmasm.translaters.macro_move_command import *
from rmasm.exceptions import *
import datetime

# =========================
#      宏语法糖命令转换器
# =========================


# key 是宏的名字
# 参数是 宏所属者 和 参数列表(list)
flattenable_macro_command_translators: Dict['str', Callable[[str, List[str]], str]] = {}


def flattenable_macro_command_translator(name):
    def wrapper(func):
        flattenable_macro_command_translators[name] = func
        return func

    return wrapper


# ================ 程序基础 ================
@flattenable_macro_command_translator('set_stack')
def f_mact_set_stack(obj: str, args: List[str]) -> str:
    if len(args) == 1:
        yield f'ss = {args[0]}'
        yield f'sp = 0'
    else:
        yield f'ss = {args[0]}'
        yield f'sp = {args[1]}'


@flattenable_macro_command_translator('push')
def f_mact_push(obj: str, args: List[str]) -> str:
    args = args.copy()
    if obj:
        args.insert(0, obj)
    yield f'sp = sp + 1'
    yield f'ss[sp] = {args[0]}'


# 调用过程，可以带条件和参数
@flattenable_macro_command_translator('call')
def f_mact_call(obj: str, args: List[str]) -> str:
    args = args.copy()
    if obj:
        args.insert(0, obj)
    cs = args[0].split()
    target = cs[0]

    call_args = []
    if len(args) > 1:
        call_args = args[1:]

    # 普通的 call
    if len(cs) == 1:
        for k, i in enumerate(call_args):
            yield f'arg{k} = {i}'
        yield f'args_len = {len(call_args)}'
        yield f'__temp = @counter + 3'
        yield f'push(__temp)'
        yield f'jump({target})'
        return

    # 下面是 if call

    if call_args:
        args_str = ',' + (','.join(call_args))
    else:
        args_str = ''
    r = random_name()

    if len(cs) > 2 and cs[1] == 'if':
        yield f'jump(call_if_{r} if {" ".join(cs[2:])})'
        yield f'jump call_end_{r}'
        yield f'call_if_{r}:'
        yield f'call({target}{args_str})'
        yield f'call_end_{r}:'
        return
    raise Exception("call 宏语句错误")


# 结束过程
@flattenable_macro_command_translator('ret')
def f_mact_ret(obj: str, args: List[str]) -> str:
    yield f'__temp = pop()'
    yield f'@counter = __temp'


# jump 的宏替换版本
@flattenable_macro_command_translator('jump')
def f_mact_jump(obj: str, args: List[str]) -> str:
    cs = args[0].split()
    csl = len(cs)
    # jump(tod if xxx)
    if csl == 1:
        yield f'jump {cs[0]}'
        return
    if cs[1] != 'if':
        raise Exception("jump 宏参数错误")

    if csl == 3:
        yield f'jump {cs[0]} if {cs[2]} == true'
        return
    if csl == 4 and cs[2] == 'not':
        yield f'jump {cs[0]} if {cs[3]} not true'
        return
    if csl == 5:
        yield f'jump {cs[0]} if {cs[2]} {cs[3]} {cs[4]}'
        return
    raise Exception("jump 宏参数错误")


@flattenable_macro_command_translator('if')
def f_mact_if(obj: str, args: List[Union[str, Code]]) -> str:
    if len(args) < 2:
        raise Exception("if 语法错误！")
    if not isinstance(args[1], Code):
        raise Exception("if 的第二个参数应该是段代码。")

    r = random_name()
    yield f'jump(if_then_{r} if {args[0]})'

    # else 结构
    if len(args) == 3:
        if not isinstance(args[2], Code):
            raise Exception("if 的第三个参数应该是段代码。")
        yield args[2]

    yield f'jump if_end_{r}'
    yield f'if_then_{r}:'
    yield args[1]
    yield f'if_end_{r}:'


@flattenable_macro_command_translator('case')
def f_mact_case(obj: str, args: List[Union[str, Code]]) -> str:
    if len(args) % 2 != 0:
        raise Exception("case 语法错误，参数数量必须是偶数，一个条件跟着一段代码。")
    k = 0
    r = random_name()
    ops = []
    codes = []
    while k < len(args):
        op = args[k]
        code = args[k + 1]
        k += 2
        if not isinstance(op, str) or not isinstance(code, Code):
            print(op.str(), '=-==-==--==---=', code.str())
            raise Exception("case 语法错误，参数必须是一个条件跟着一段代码。")
        ops.append(op)
        codes.append(code)

    for k, op in enumerate(ops):
        yield f'jump(case_then_{k}_{r} if {op})'
    yield f'jump(case_end_{r})'
    for k, code in enumerate(codes):
        yield f'case_then_{k}_{r}:'
        yield code
        if k < len(codes) - 1:
            yield f'jump(case_end_{r})'
    yield f'case_end_{r}:'


@flattenable_macro_command_translator('while')
def f_mact_while(obj: str, args: List[str]) -> str:
    if len(args) != 2:
        raise Exception("while 语法错误！")
    if not isinstance(args[1], Code):
        raise Exception("while 第二个参数应该是段代码。")
    r = random_name()
    yield f'jump(while_start_{r} if {args[0]})'
    yield f'jump while_end_{r}'
    yield f'while_start_{r}:'
    yield args[1]
    yield f'jump(while_start_{r} if {args[0]})'
    yield f'while_end_{r}:'


# 导入文件
@flattenable_macro_command_translator('include')
def f_mact_include(obj: str, args: List[str]) -> str:
    file_path = args[0]
    if file_path.startswith('"') and file_path.endswith('"'):
        file_path = file_path[1:-1]
    else:
        raise Exception("导入的文件路径必须是字符串形式")

    try:
        fp = open(file_path, 'r', encoding='utf-8')
        f = fp.read()
        fp.close()
    except Exception as e:
        raise Exception(f"打开文件 {file_path} 失败，可能是文件不存在或编码不是 utf-8。")

    yield Code(content=f)


macro_fragment: Dict[str, Code] = {}


@flattenable_macro_command_translator('define')
def f_mact_define(obj: str, args: List[Union[str, Code]]) -> str:
    args = args.copy()
    if obj:
        args.insert(0, obj)

    if not isinstance(args[1], Code):
        raise Exception("宏定义的语法错误，第二个参数应该是段代码。")

    macro_fragment[args[0]] = args[1]
    yield ''


@flattenable_macro_command_translator('macro')
def f_mact_macro(obj: str, args: List[str]) -> str:
    args = args.copy()
    if obj:
        args.insert(0, obj)
    name = args[0]
    del args[0]
    if name not in macro_fragment:
        raise Exception(f"要是用的宏不存在：{name}")
    source_code = macro_fragment[name]

    # 宏替换参数
    kv_map = {}
    for k, v in enumerate(args):
        if isinstance(v, Code):
            kv_map[f'arg{k}'] = v.content
        else:
            kv_map[f'arg{k}'] = v

    # 宏替换额外参数
    kv_map['args_len'] = len(args)
    kv_map['random_name'] = random_name()
    kv_map['macro_name'] = name
    kv_map['time'] = str(datetime.datetime.now())

    new_code = Code(content=code_replace(source_code.content, kv_map))
    yield new_code


# ================ 快捷语法 ================

# 内存到内存
@flattenable_macro_command_translator('mtom')
@flattenable_macro_command_translator('read_write')
def f_mact_read_write(obj: str, args: List[str]) -> str:
    yield f'__temp = {args[0]}'
    yield f'{args[1]} = __temp'


# 自增
@flattenable_macro_command_translator('add')
def f_mact_add(obj: str, args: List[str]) -> str:
    if len(args) == 1:
        args = args.copy()
        args.append('1')
    yield f'{args[0]} = {args[0]} + {args[1]}'


# 自减
@flattenable_macro_command_translator('sub')
def f_mact_sub(obj: str, args: List[str]) -> str:
    if len(args) == 1:
        args = args.copy()
        args.append('1')
    yield f'{args[0]} = {args[0]} - {args[1]}'


# 快速获取属性
@flattenable_macro_command_translator('get')
def f_mact_get(obj: str, args: List[str]) -> str:
    # 'ww,dd,ee,rr'
    for arg in args:
        yield f'{obj}-{arg} = {obj}.@{arg}'


# ================ 设备操作 ================

@flattenable_macro_command_translator('show')
def f_mact_println(obj: str, args: List[str]) -> str:
    yield f'print({args[0]})'
    yield f'print_flush({obj})'
